import torch
import torch as t
import matplotlib.pyplot as plt

# 为CPU中设置种子，生成随机数：
t.manual_seed(100)
dtype = t.float

# 生成x坐标数据，x为tensor,需要把x的形状转换为100x1
x = t.unsqueeze(t.linspace(-1, 1, 100), dim=1)

# 生成y的坐标数据，y为tensor，形状为100x1，另加上一些噪声
y = 3 * x.pow(2) + 2 + 0.2 * torch.rand(x.size())

# 画图，把tensor数据转为numpy数据
# plt.scatter(x.numpy(), y.numpy())
# plt.show()

# 初始化权重参数
w = t.randn(1, 1, dtype=dtype, requires_grad=True)
b = t.zeros(1, 1, dtype=dtype, requires_grad=True)

# 训练模型

lr = 0.001  # 学习率

for ii in range(800):
    # 前向传播，并定义损失函数loss
    y_pred = x.pow(2).mm(w) + b
    loss = 0.5 * (y_pred - y) ** 2
    loss = loss.sum()

    # 自动计算梯度，梯度放在grad属性中
    loss.backward()

    with t.no_grad():
        w -= lr * w.grad
        b -= lr * b.grad

    w.grad.zero_()
    b.grad.zero_()

# 可视化训练结果
plt.plot(x.numpy(), y_pred.detach().numpy(), 'r--', label='predict')
plt.scatter(x.numpy(), y_pred.detach().numpy(), color='b', label='true')

plt.xlim(-1, 1)
plt.ylim(2, 6)
plt.legend()
plt.show()

print(w, b)
